import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import cv2
from collections import defaultdict
from pathlib import Path
from PIL import Image

# 为兼容源码进行train_df数据预处理
first_df = pd.read_csv("../data/steel_data/train.csv")
train_df = pd.DataFrame()
result = []
# 根据每个ImageId的出现次数，将每个ImageId对应的数据凑足四行。对于不足四行的ImageId，需要补充数据行，其中ImageId列复制，其他列设置为空
for _, group in first_df.groupby('ImageId'):
    num_missing_rows = 4 - len(group)
    if num_missing_rows > 0:
        missing_rows = pd.DataFrame({
            'ImageId': [group['ImageId'].iloc[0]] * num_missing_rows,
            'ClassId': [np.nan] * num_missing_rows,
            'EncodedPixels': [np.nan] * num_missing_rows
        })
        group = pd.concat([group, missing_rows], ignore_index=True)
    result.append(group.head(4))
# 根据 ClassId 的值（例如 1、2、3、4）来决定每行在同一个 ImageId 下的具体位置。例如ClassId是3.0，则移动到第3行
temp_df = pd.concat(result, ignore_index=True)
for image_id in temp_df['ImageId'].unique():
    subset = temp_df[temp_df['ImageId'] == image_id]
    new_subset = [None] * 4
    for idx, row in subset.iterrows():
        class_id = row['ClassId']
        if pd.notna(class_id):
            new_subset[int(class_id) - 1] = row
        else:  # 保持原位
            for i in range(4):
                if new_subset[i] is None:
                    new_subset[i] = row
                    break
    train_df = pd.concat([train_df, pd.DataFrame(new_subset)], ignore_index=True)  # 完成train_df预处理
# print(train_df.head(40))
train_size_dict = defaultdict(int)
train_path = Path("../data/steel_data/train_images/")
for img_name in train_path.iterdir():
    img = Image.open(img_name)
    train_size_dict[img.size] += 1

# 训练集中图像的尺寸和数目
print(train_size_dict)

# 读取测试集图像数据
test_size_dict = defaultdict(int)
test_path = Path("../data/steel_data/test_images/")

for img_name in test_path.iterdir():
    img = Image.open(img_name)
    test_size_dict[img.size] += 1
print(test_size_dict)

# 为不同的缺陷类别设置颜色显示
palet = [(249, 192, 12), (0, 185, 241), (114, 0, 218), (249, 50, 12)]
fig, ax = plt.subplots(1, 4, figsize=(15, 5))
for i in range(4):
    ax[i].axis('off')
    ax[i].imshow(np.ones((50, 50, 3), dtype=np.uint8) * palet[i])
    ax[i].set_title("class color: {}".format(i + 1))
fig.suptitle("each class colors")
plt.show()

# 将不同的缺陷标识归类
idx_class_1 = []
idx_class_2 = []
idx_class_3 = []
idx_class_4 = []
idx_class_multi = []

for col in range(0, len(train_df), 4):
    img_names = [str(i).split("_")[0] for i in train_df.iloc[col:col + 4, 0].values]
    if not (img_names[0] == img_names[1] == img_names[2] == img_names[3]):
        raise ValueError
    labels = train_df.iloc[col:col + 4, 1].values.tolist()
    # 完全分类
    # if labels[0] == 1:
    #     idx_class_1.append(col)
    # if labels[1] == 2:
    #     idx_class_2.append(col)
    # if labels[2] == 3:
    #     idx_class_3.append(col)
    # if labels[3] == 4:
    #     idx_class_4.append(col)
    # 单独与多种故障分开
    nan_count = sum(np.isnan(label) for label in labels)
    if nan_count == 3:
        if labels[0] == 1:
            idx_class_1.append(col)
        elif labels[1] == 2:
            idx_class_2.append(col)
        elif labels[2] == 3:
            idx_class_3.append(col)
        elif labels[3] == 4:
            idx_class_4.append(col)
    else:
        idx_class_multi.append(col)

train_path = Path("../data/steel_data/train.csv")


# 可视化标注函数
def name_and_mask(start_idx):
    col = start_idx
    img_names = [str(i) for i in train_df.iloc[col:col + 4, 0].values]
    label_s = [i for i in train_df.iloc[col:col + 4, 1].values]
    if not (img_names[0] == img_names[1] == img_names[2] == img_names[3]):
        raise ValueError
    labels = train_df.iloc[col:col + 4, 2]
    mask = np.zeros((256, 1600, 4), dtype=np.uint8)
    for idx, label in enumerate(labels.values):
        if label is not np.nan:
            mask_label = np.zeros(1600 * 256, dtype=np.uint8)
            label = label.split(" ")
            positions = map(int, label[0::2])  # 以2为步长跳跃
            length = map(int, label[1::2])
            for pos, le in zip(positions, length):
                mask_label[pos - 1:pos + le - 1] = 1
            mask[:, :, idx] = mask_label.reshape(256, 1600, order='F')  # 按列取值reshape
    return label_s, img_names[0], mask


def save_mask_image(col, class_id):
    lable_s, name, mask = name_and_mask(col)
    img = cv2.imread(str(train_path / name))
    fig, ax = plt.subplots(figsize=(16, 2.56))
    for ch in range(4):
        contours, _ = cv2.findContours(mask[:, :, ch], cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
        for i in range(0, len(contours)):
            cv2.polylines(img, contours[i], True, palet[ch], 2)
            print(contours[i])

    # ax.imshow(img)
    # ax.set_xticks([])
    # ax.set_yticks([])
    # plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
    # plt.margins(0, 0)
    # plt.savefig(str(Path("./data/after/" + class_id) / name))


for idx in idx_class_1:  # 第1类缺陷
    save_mask_image(idx, "class_1")
for idx in idx_class_2:  # 第2类缺陷
    save_mask_image(idx, "class_2")
for idx in idx_class_3:  # 第3类缺陷
    save_mask_image(idx, "class_3")
for idx in idx_class_4:  # 第4类缺陷
    save_mask_image(idx, "class_4")
for idx in idx_class_multi[:1]:  # 同时具有多种缺陷
    save_mask_image(idx, "class_multiple")
